# -*- coding: utf-8 -*-

from __future__ import print_function
import sys,os
sys.path.append('/opt/alps/lib') # for pyalps
os.environ['PATH'] = os.environ['PATH'].replace('/share/opt/alps/bin','/opt/alps/bin') # for mps_optim etc
os.environ['OMP_NUM_THREADS'] = '6' # cores for OpenMP in single task
import pyalps
import matplotlib.pyplot as plt
import numpy as np
import time
import glob, shutil
import multiprocessing
from numpy import array,shape,arange
from data_analysis_functions import get_obervable_v4,get_correlation_matrices_v4,get_truncated_weight
from data_analysis_functions import gather_2D_correlation_data
from generate_ssh_ladder import generate_ssh_ladder_lattice

'''
Boson SSH ladder, theory gauge
'''

#----------------------1 Run the Task------------------------

run_with_GUI = 0
numcores = 1
run_the_task = 1
pi = np.pi
t1 = 1.0
t2 = 0.3
K = 1.0
#phis = [0,0.1*pi,0.3*pi,0.5*pi,0.65*pi,0.7*pi,0.9*pi,0.95*pi,pi,1.1*pi]
#phis = [0.25*pi,0.3*pi,0.75*pi,0.9*pi,pi]
phis_over_pi = np.linspace(0,1,21)
phi_over_pi = phis_over_pi[-1]
#Ds = [200,500,800,1000,1200,1500,2000]
#for K in arange(0,3.6,0.5):
#    for phi_over_pi in arange(0,1.1,0.2):
#        Paras.append((K,phi_over_pi))
#Paras = [600]
task_name = 'ssh_ladder_for_boson' 
sweeps = 20
maxstates = 800
#L = 50
Paras = []
for L in [50]:
    (lattice_library,lattice_name)=generate_ssh_ladder_lattice(L) # OBC ladder
    Paras.append((lattice_library,lattice_name,L))
#Paras = Ds
n = 0.1
N_total = int(n*L*2)
#phi = 0.25*pi

# bad_ids works only if supplement==1, supplemnt works only when run_the_task==1
bad_ids = []
supplement = 0
sup_sweeps = 20
sup_maxstates = 1000
if supplement==1:
    Paras = Paras[bad_ids]
    sweeps = sup_sweeps
    maxstates = sup_maxstates
os.environ["OMP_NUM_THREADS"] = '6' 

def Runmps(para):
#    (K,phi_over_pi) = para
    (lattice_library,lattice_name,L) = para
    para = L
    phi = phi_over_pi*pi
    ##infomation hidden,ie (maxstates,sweeps) = para, or sweeps=para
    parm = {}
    
    parm['LATTICE_LIBRARY'] = lattice_library
    parm['LATTICE'] = lattice_name
    parm['L'] = L
    parm['MODEL_LIBRARY'] = 'boson_ssh_ladder.xml'
    parm['MODEL'] = 'boson_ssh_ladder'
    parm['COMPLEX'] = 1
    
    parm['t1'] = t1
    parm['t2'] = t2
    parm['K'] = K
    parm['phi'] = phi
    parm['U'] = 0 # 0 means no interaction
    parm['Nmax'] = 1 # 1 for hardcore boson, then U is meaningless
    parm['N_total'] = N_total
    parm['CONSERVED_QUANTUMNUMBERS'] = 'N'
    parm['MAXSTATES'] = maxstates
    parm['SWEEPS'] = sweeps
    parm['ietl_jcd_maxiter'] = 20
#    parm['TRUNCATION'] = 1e-8
    parm['NUMBER_EIGENVALUES'] = 1
    parm['MEASURE[EnergyVariance]'] = 1
    parm['MEASURE[Entropy]'] = 1
    
    parm['MEASURE_LOCAL[local_density]'] = 'n'
    parm['MEASURE_LOCAL[local_density_square]'] = 'ni2'

#    parm['MEASURE_CORRELATIONS[curr_correlation]'] = 'bdag:b'
#    parm['MEASURE_LOCAL[site_operator1]'] = 'site_operator1'
    
    list1 = [] # for t1*exp(-1i*phi/2) L odd
    list2 = [] # for t1*exp(1i*phi/2) R odd
    list3 = [] # for t2*exp(-1i*phi/2) L even
    list4 = [] # for t2*exp(1i*phi/2) R even
    ID = 1
    for i in range(1,2*L-2,2):
        if ID==1:
            list1.append(i)
            list2.append(i+1)
            ID = 2
        elif ID==2:
            list3.append(i)
            list4.append(i+1)
            ID = 1
    # here vertex has to be from 0
    list1 = [i-1 for i in list1]
    list2 = [i-1 for i in list2]
    list3 = [i-1 for i in list3]
    list4 = [i-1 for i in list4]
    
    Ledges_odd = ','.join([str((i,i+2)) for i in list1])
    Redges_odd = ','.join([str((i,i+2)) for i in list2])
    Ledges_even = ','.join([str((i,i+2)) for i in list3])
    Redges_even = ','.join([str((i,i+2)) for i in list4])
    
    list5 = [2*i-1 for i in range(1,L+1)] 
    inter_edges = ','.join([str((i-1,i)) for i in list5]) # for [1 2],[3 4]..., L/R or inter legs

#    measured_pair = [str((i,i+1)) for i in range(L-1)] ## 待会看关联的空间分布
#    measured_pair = ','.join(measured_pair)
#    parm['MEASURE_LOCAL_AT[up_correlation]'] = 'bdag:b|'+measured_pair
    
    parm['MEASURE_LOCAL_AT[L_correlations_odd]'] = 'bdag:b|'+Ledges_odd # vertex ID-1
    parm['MEASURE_LOCAL_AT[R_correlations_odd]'] = 'bdag:b|'+Redges_odd # vertex ID-1
    parm['MEASURE_LOCAL_AT[L_correlations_even]'] = 'bdag:b|'+Ledges_even # vertex ID-1
    parm['MEASURE_LOCAL_AT[R_correlations_even]'] = 'bdag:b|'+Redges_even # vertex ID-1
    parm['MEASURE_LOCAL_AT[inter_correlations]'] = 'bdag:b|'+inter_edges # vertex ID-1
    
    parm = [parm]

    input_file = pyalps.writeInputFiles(task_name + '_para=' + str(para) + '_task', parm)
    time1 = time.time()
    pyalps.runApplication('mps_optim', input_file, writexml=True)
    time2 = time.time()
    #delete_MPS_wavefunction(task_name + '_para=' + str(para) + '_task')  # 注释此句来决定是否删除chkp文件夹 #子函数不分先后
    if run_with_GUI == 0:
        doc1 = open('1Run_print.txt', 'a')
        print(
            'para = %s task finished with time %.2f minutes.' %
            (para, (time2 - time1) / 60.0),
            file=doc1)
        doc1.close()
    else:
        print('para = %s task finished with time %.2f minutes.' %
              (para, (time2 - time1) / 60.0))

def delete_old_outfiles():  # 删除上一次运行得到的文件
    old_in_files = glob.glob('*%s*.in.*' % task_name)
    old_out_files = glob.glob('*%s*.out.*' % task_name)
    for i in old_in_files + old_out_files:
        try:
            os.remove(i)  #不能删文件夹
        except OSError:
            try:
                shutil.rmtree(i)  #用来删文件夹
            except:
                print('Something wrong when deleting files.')
##如果不需要接着计算，删除占用大的MPS波函数chkp文件夹,就运行这个函数
def delete_MPS_wavefunction(single_task_string):
    chkp_files = glob.glob('*%s*out*chkp' % single_task_string)
    for i in chkp_files:
        try:
            shutil.rmtree(i)  #用来删文件夹
        except:
            print('Maybe no such file to delete.')
    print('check point files deleted.')

if run_the_task == 1:
    pool = multiprocessing.Pool(processes=numcores)
    ##    Paras = [i+1 for i in range(20)]
    print('report: This task has %d parallel task/tasks.' % (len(Paras)))
    print('task %s is Running...' % (task_name))
    pool.map(Runmps, Paras, 1)
    pool.close()
    pool.join()
    if run_with_GUI == 0:
        doc1 = open('1Run_print.txt', 'a')
        print('Finished with all tasks.', file=doc1)
        doc1.close()
    else:
        print('Finished with all tasks.')
else:
    print('Task did not run, continue with data analysis.')
## 如果run_the_task为0，就只有最开始的一点会运行

#----------------------2 结果分析----------------------------------------------
result_files = pyalps.getResultFiles(prefix=task_name)
assert result_files != [], 'Report: Error, no result file has been calculated out. The calculation has definitly failed.'
## 如果没有产生结果文件，马上报错，而不是隐藏着继续运行
if run_with_GUI == 0:
    plt.switch_backend('pdf')
## 解决无GUI界面不能正常画图的问题

### The usage examples are as follows:
(Ls, energy_variance) = get_obervable_v4(result_files,'L', 'EnergyVariance')
#print(Ds)
print('energy_variance is %s'%(energy_variance))
(Ls, truncation_error) = get_truncated_weight(result_files,'L')
print('truncation error is %s'%(truncation_error))
print('corresponding Ls are %s'%(Ls))
#ids = arange(len(phis))
#bad_VarH_ids = ids[energy_variance>1e-8]
#print('bad VarH tasks are %s (from 0)'%(bad_VarH_ids))

doc = open('2Accuracy_evaluation.txt', 'w')
print('corresponding Ls are %s'%(Ls),file=doc)
print('energy_variance is %s'%(energy_variance),file=doc)
print('truncation error is %s'%(truncation_error),file=doc)
#print('bad VarH tasks are %s (from 0)'%(bad_VarH_ids),file=doc)
doc.close()

### does Eg bears kink?
#(phis, Egs) = get_obervable_v4(result_files,'phi', 'Energy')
#Egs_derivative = np.diff(Egs)/np.diff(phis)
### 画图
#fig = plt.figure()
#(fig, ax) = plt.subplots(figsize=(4, 3), dpi=120)
#ax.plot(phis[0:-1]/pi, Egs_derivative)
##ax.legend(loc=1,fontsize=8)
#ax.set_xlabel('phi(pi)')
#ax.set_ylabel('Energy Derivative')
##ax.set_ylim([1e-8,1e2])
##ax.set_title('SPINmodel MPS convergence')
#plt.tight_layout() ##solve figure in pdf saved cutted off
#fig.savefig('Energy_Derivative_with_phi.pdf')
##import scipy.io
##scipy.io.savemat('data.mat',{'phi':phis,'Egs':Egs,\
##                             'Ds':Ds,'energy_variance':energy_variance,\
##                             'truncation_error':truncation_error})


## fit Von Neumann entropy data to get central charge. The ladder is OBC
(Ls, entropys_phi) = get_obervable_v4(result_files,'L', 'Entropy')
# which_phi = 2-1
for which_L in range(len(Ls)):
    entropys = entropys_phi[which_L]
    L = Ls[which_L]
    rung_cuts = arange(2,L*2,2)
    ladder_entropys = entropys[rung_cuts-1]
    cuts = rung_cuts/2

    ### 拟合entropy curve naive
    #from scipy.optimize import curve_fit
    #def fit_func1(j,c,B):
    #    inside = (L/pi)*np.sin(pi*j/L)
    #    return c/6.0*np.log(inside)+B
    #js = arange(1,len(ladder_entropys)+1)
    #(popt, pcov) = curve_fit(fit_func1, js, ladder_entropys)
    #(c,B) = popt
    ##print('c is %f'%(c))

    ## correlations
    ## scan phi
    def get_1D_scan_correlations(propstring,observable):
        eigen_measure = pyalps.loadEigenstateMeasurements(result_files,observable)
        eigen_measure = [i[0] for i in eigen_measure]
        def by_props(dataset):  ##把DataSet按照其props中的mu值排序
                return dataset.props[propstring]
        obs_sorted = sorted(eigen_measure, key=by_props)
        props = array([i.props[propstring] for i in obs_sorted])
        correlations = array([i.y[0] for i in obs_sorted])
        return (props,correlations)

    (Ls,L_correlations_odd) = get_1D_scan_correlations('L','L_correlations_odd')
    (Ls,R_correlations_odd) = get_1D_scan_correlations('L','R_correlations_odd')
    (Ls,L_correlations_even) = get_1D_scan_correlations('L','L_correlations_even')
    (Ls,R_correlations_even) = get_1D_scan_correlations('L','R_correlations_even')
    (Ls,inter_correlations) = get_1D_scan_correlations('L','inter_correlations')

    L_correlations_odd = L_correlations_odd[which_L]
    R_correlations_odd = R_correlations_odd[which_L]
    L_correlations_even = L_correlations_even[which_L]
    R_correlations_even = R_correlations_even[which_L]
    inter_correlations = inter_correlations[which_L]

    #####################################
    ### 正确的hopping energy 相位是不能丢的,顺便把t1 t2弄进去吧,弄最合理的hopping energy
#    phi = phis[which_phi]
    phi = phi_over_pi*pi
    L_hopping_energy_odd = 2*np.real(t1*np.exp(-1j*phi/2)*L_correlations_odd) # z+z*=2Re(z)
    R_hopping_energy_odd = 2*np.real(t1*np.exp(1j*phi/2)*R_correlations_odd)
    L_hopping_energy_even = 2*np.real(t2*np.exp(-1j*phi/2)*L_correlations_even)
    R_hopping_energy_even = 2*np.real(t2*np.exp(1j*phi/2)*R_correlations_even)
    #inter_hopping_energy = 2*np.real(inter_correlations)
    #####################################

    def cross(a1,a2):
        (L1,L2) = (len(a1),len(a2))
        L = L1+L2
        mL = min(L1,L2)
        if L1>L2:
            remain = a1[L2:]
            a1 = a1[:L2]
        elif L1<L2:
            remain = a2[L1:]
            a2 = a2[:L1]
        else:
            remain = []
        b = np.zeros(L)
        b[:] = np.NaN
        # a1 first
        b[2*arange(mL)] = a1
        b[2*arange(mL)+1] = a2  
        b[2*mL:] = remain
        return b

    L_hopping_energy = cross(L_hopping_energy_odd,L_hopping_energy_even)
    R_hopping_energy = cross(R_hopping_energy_odd,R_hopping_energy_even)
    hopping_energy = L_hopping_energy+R_hopping_energy

    ### 拟合entropy curve with orscillation
    from scipy.optimize import curve_fit
    from scipy.interpolate import interp1d

    js = cuts
    hopping_energy_func = interp1d(js,hopping_energy,kind='cubic')

    def fit_func2(j,c,A,B):
        inside = (L/pi)*np.sin(pi*j/L)
        return c/6.0*np.log(inside)+A+B*hopping_energy_func(j)
    (fit_para, pcov) = curve_fit(fit_func2, js, ladder_entropys)
    (c,A,B) = fit_para

#    # L_hopping_energy R_hopping_energy should be continus though
#    # hopping energy fit
#    fig = plt.figure()
#    (fig, ax) = plt.subplots(figsize=(4, 3), dpi=120)
#    ax.plot(js, hopping_energy,'.')
#    js_dense = arange(1,L-1,0.1)
#    js_dense = np.append(js_dense,L-1)
#    ax.plot(js_dense, hopping_energy_func(js_dense),'-')
#    #ax.set_ylim([-0.01,0.01])
#    ax.set_xlabel('L cut')
#    ax.set_ylabel('L or R hopping energy')

    ### interleg hopping energy test
    #js_full = arange(1,L+1)
    #fig = plt.figure()
    #(fig, ax) = plt.subplots(figsize=(4, 3), dpi=120)
    #ax.plot(js_full, inter_hopping_energy,'.-')
    #ax.set_xlabel('j')
    #ax.set_ylabel('interleg hopping energy')

    ## 画图
    # entropy fig
    fig = plt.figure()
    (fig, ax) = plt.subplots(figsize=(4, 3), dpi=120)
    ax.plot(cuts, ladder_entropys,'.')
    js_dense = arange(1,L-1,0.1)
    js_dense = np.append(js_dense,L-1)
    ax.plot(js_dense, fit_func2(js_dense,c,A,B))
    #ax.legend(loc=1,fontsize=8)
    ax.set_xlabel('L cut')
    ax.set_ylabel('entropy')
    #ax.set_ylim([1e-8,1e2])
    ax.set_title('c is %f'%(c))
    plt.tight_layout() ##solve figure in pdf saved cutted off
    fig.savefig('central_charge_with_L_%s.pdf'%(Ls[which_L]))
    #import scipy.io
    #scipy.io.savemat('data.mat',{'phi':phis,'Egs':Egs,\
    #                             'Ds':Ds,'energy_variance':energy_variance,\
    #                             'truncation_error':truncation_error})


## 画图代码备用
#fig = plt.figure()
#(fig, ax) = plt.subplots(figsize=(4, 3), dpi=300)
#for i in range(len(result_files)): # 多个curve自动上色
#    ax.plot(xdatas[i], ydatas[i]-ydatas[i][-1], label='L=%s'%(label_list[i]))
#ax.legend(loc=1,fontsize=8)
#ax.set_xlabel('MPS sweeps')
#ax.set_ylabel('Energy Difference with final')
#ax.set_yscale('log')
##ax.set_ylim([1e-8,1e2])
#ax.grid()
#ax.set_title('SPINmodel MPS convergence')
#plt.tight_layout() ##solve figure in pdf saved cutted off
#fig.savefig('SPINmdoel_sweeps_vs_energy.pdf')
#import scipy.io
#scipy.io.savemat('data.mat',{'phi':phis,'K':Ks,'Z':Z})
